#!/usr/bin/env python
# external_command: Thread-safe execution of external processes
#
# Author:   Mike Murphy (mamurph@cs.clemson.edu)
# Revision: 26 February 2009
#
#   Copyright 2007-2009 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.


# NOTE: This code probably only works on Linux and Mac systems

# Read buffer size and fallback for MAXFD
READ_SIZE = 16384
MAXFD_FALLBACK = 1024

import os, resource, select, sys, threading


# SignalMonitor is a monitor object whose sole purpose in life is to store
# an integer (in a thread-safe manner), so that subprocess management can
# poll the monitor to retrieve any signals that might have been sent.
class SignalMonitor(object):
   def __init__(self):
      self.sema = threading.BoundedSemaphore(1)
      self.signal = 0
   def send_signal(self, signum):
      self.sema.acquire()
      self.signal = signum
      self.sema.release()
   def check_signal(self):
      signum = 0
      self.sema.acquire()
      signum = self.signal
      self.sema.release()
      return signum
#


# A thread-safe way to run external programs from Python, with support for
# retrieving the exit code, receiving stdout and stderr data, sending data
# to subprocess stdin, and propagating signals from the parent process to
# the child(ren)
def run_external(command, args, stdin='', sigmon = None):
   stdout = ''
   stderr = ''
   in_idx = 0
   
   # Make command args[0]
   args.insert(0, command)
   
   # Define pipes for I/O
   in_pipe_r, in_pipe_w = os.pipe()
   out_pipe_r, out_pipe_w = os.pipe()
   err_pipe_r, err_pipe_w = os.pipe()
   
   # Need a flag to determine if the write is done. Polling a closed fd in
   # select results in an error.
   write_done = False
   
   # Fork the child (which has pid 0)
   pid = os.fork()
   if pid == 0:
      # IMPORTANT: The write end of the input pipe MUST be closed in the child,
      # or EOF will not be sent when the write end is closed in the parent.
      # For security, go ahead and close all open fd's EXCEPT in_pipe_r,
      # out_pipe_w, and err_pipe_w
      maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
      if (maxfd == resource.RLIM_INFINITY):
         maxfd = MAXFD_FALLBACK
      for fd in range(0, maxfd):
         try:
            if fd not in (in_pipe_r, out_pipe_w, err_pipe_w):
               os.close(fd)
         except OSError:
            pass
      
      # dup the remaining fd's onto stdin, stdout, and stderr, then exec
      # the child operation
      os.dup2(in_pipe_r, 0)
      os.dup2(out_pipe_w, 1)
      os.dup2(err_pipe_w, 2)
      os.execvp(command, args)
      # BYE: child exits Python with the exec
   ######################################################################
   
   # MUST be in parent here: exec in child would have exited Python
   
   # Do a prime read above the loop, in case a short-running child
   # terminates before we get through all the polling. To be safe, all
   # further os.waitpid calls after this one should be conditional on
   # (opid, status) == (0, 0)
   opid, status = os.waitpid(pid, os.WNOHANG)
   out_avail, in_avail, throwaway = \
                 select.select([out_pipe_r, err_pipe_r], [], [], 0)
   
   # Loop condition: child process has not terminated, or all available
   # child-produced data (stdout and stderr) has not yet been retrieved
   while (opid, status) == (0, 0) or len(out_avail) > 0:
      # If a closed fd is presented to select, an exception will result
      in_sel = []
      if not write_done:
         in_sel = [in_pipe_w]
      
      # select semantics follow those of select(2) in C
      out_avail, in_avail, throwaway = \
                 select.select([out_pipe_r, err_pipe_r], in_sel, [], 0)
      
      # Intermittent poll on the child, in case the OS cleans up the results
      # while we are busy trying to collect data
      if (opid, status) == (0, 0):
         opid, status = os.waitpid(pid, os.WNOHANG)
      
      # Standard input *to* the subprocess. This data needs to be sent using
      # the async select mechanism, in case the data to be sent exceeds the
      # buffer provided by the OS.
      if len(in_avail) > 0:
         in_idx = os.write(in_avail[0], stdin[in_idx:])
         if in_idx >= len(stdin):
            os.close(in_pipe_w)
            write_done = True
      
      # Another intermittent poll
      if (opid, status) == (0, 0):
         opid, status = os.waitpid(pid, os.WNOHANG)
      
      # Standard output and error *from* the subprocess
      for fd in out_avail:
         data = os.read(fd, READ_SIZE)
         if fd == out_pipe_r:
            stdout += data
         else:
            stderr += data
         # Poll between fd's, for good measure
         if (opid, status) == (0, 0):
            opid, status = os.waitpid(pid, os.WNOHANG)
      
      # Check for signals to send to the child
      if sigmon:
         signum = sigmon.check_signal()
         if signum != 0:
            try:
               os.kill(pid, signum)
            except Exception:
               pass
      
      # Poll for child process status one last time this iteration.
      if (opid, status) == (0, 0):
         opid, status = os.waitpid(pid, os.WNOHANG)
      
      # If the subprocess has finished, do one last check to be sure no
      # output is lying around waiting to be read
      if (opid, status) != (0, 0):
         out_avail, in_avail, throwaway = \
                 select.select([out_pipe_r, err_pipe_r], [], [], 0)
   #
   
   # Done with the child process: handle the results
   sig = status & 0x007f        # Strip high bit: don't care about core dump
   code = (status & 0xff00) >> 8
   if sig != 0:
      code = 256 - sig
   return (code, stdout, stderr)
#



if __name__ == '__main__':
   stdin = open(sys.argv[1], 'r').read()
   sigmon = SignalMonitor()
   #sigmon.send_signal(2)
   code, out, err = run_external('/bin/cat', [], stdin, sigmon)
   print out
   print '*** (1)', code
   
   # Short-running test case
   code, out, err = run_external('/bin/false', [], sigmon=sigmon)
   print '*** (2)', code
